from flask import Flask, render_template, Response, request, jsonify
import cv2
import logging
import numpy as np
import os
from flask_cors import CORS
import requests
from utils import init_model, process_frame

app = Flask(__name__)
CORS(app)  # 允许所有域名的请求
logging.basicConfig(format='%(asctime)s:%(levelname)s:%(message)s', level=logging.INFO)

UPLOAD_FOLDER = 'uploads'
UPLOAD_URL = "http://localhost:5000/upload_video"
VIDEO_FEED_URL = "http://localhost:5000/video_feed/2"

if not os.path.exists(UPLOAD_FOLDER):
    os.makedirs(UPLOAD_FOLDER)

model = None
model2 = None
model2_name = None
recognition_status = False
cap = None
frame_position = 0
is_paused = False

SAFE_DISTANCE = 50  # 安全距离（像素）
region = [100, 100, 400, 400]  # 默认区域

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/video_feed')
def video_feed():
    return Response(generate_frames(), mimetype='multipart/x-mixed-replace; boundary=frame')

@app.route('/stream_feed')
def stream_feed():
    return Response(generate_stream_frames(VIDEO_FEED_URL), mimetype='multipart/x-mixed-replace; boundary=frame')


@app.route('/start_recognition', methods=['POST'])
def start_recognition():
    global recognition_status
    recognition_status = True  # 启动目标识别
    return jsonify({'status': 'Target recognition started'}), 200

def send_frame_to_server(frame):
    _, buffer = cv2.imencode('.jpg', frame)
    try:
        response = requests.post(UPLOAD_URL, data=buffer.tobytes(), headers={'Content-Type': 'application/octet-stream'})
        if response.status_code != 200:
            logging.error(f"Failed to upload frame: {response.status_code}")
    except Exception as e:
        logging.error(f"Exception occurred while uploading frame: {e}")

def generate_frames():
    global model, model2, recognition_status, cap, frame_position, is_paused, region
    cap = cv2.VideoCapture(VIDEO_FEED_URL)  # Use 0 for webcam or VIDEO_FEED_URL for a stream URL

    # 设置帧大小
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 480)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
    frame_skip = 5  # 跳过的帧数
    frame_count = 0

    while True:
        if is_paused:
            continue

        success, frame = cap.read()
        if not success:
            break
        else:
            frame_count += 1
            if frame_count % frame_skip == 0:
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                if recognition_status and model:
                    if model2_name == 'detection_module.pt':
                        frame = process_frame(model, model2, frame, region, True, True)
                    else:
                        frame = process_frame(model, frame, region, True, True)
                send_frame_to_server(frame)  # 上传处理后的帧
                ret, buffer = cv2.imencode('.jpg', frame, [int(cv2.IMWRITE_JPEG_QUALITY), 70])
                if not ret:
                    continue
                frame = buffer.tobytes()
                yield (b'--frame\r\n'
                       b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')
    cap.release()

@app.route('/upload_video', methods=['POST'])
def upload_video():
    if 'file' not in request.files:
        return jsonify({'error': 'No file part'}), 400
    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': 'No selected file'}), 400

    file_path = os.path.join(UPLOAD_FOLDER, file.filename)
    file.save(file_path)

    return Response(process_local_video(file_path), mimetype='multipart/x-mixed-replace; boundary=frame')

def generate_stream_frames(url):
    global model, model2, recognition_status, region
    cap = cv2.VideoCapture(url)

    # 设置帧大小
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 480)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
    frame_skip = 5  # 跳过的帧数
    frame_count = 0

    while True:
        success, frame = cap.read()
        if not success:
            break
        else:
            frame_count += 1
            if frame_count % frame_skip == 0:
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                if recognition_status and model:
                    if model2_name == 'detection_module.pt':
                        frame = process_frame(model, model2, frame, region, True, True)
                    else:
                        frame = process_frame(model, frame, region, True, True)
                send_frame_to_server(frame)  # 上传处理后的帧
                ret, buffer = cv2.imencode('.jpg', frame, [int(cv2.IMWRITE_JPEG_QUALITY), 70])
                if not ret:
                    continue
                frame = buffer.tobytes()
                yield (b'--frame\r\n'
                       b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')
    cap.release()

def process_local_video(file_path):
    global model, model2, recognition_status, cap, frame_position, is_paused
    cap = cv2.VideoCapture(file_path)

    while True:
        if is_paused:
            continue

        success, frame = cap.read()
        if not success:
            break
        else:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            if recognition_status and model:
                if model2_name == 'detection_module.pt':
                    frame = process_frame(model, model2, frame,region, True, True)
                else:
                    frame = process_frame(model, frame,region, True, True)
            send_frame_to_server(frame)  # 上传处理后的帧
            _, buffer = cv2.imencode('.jpg', frame)
            frame = buffer.tobytes()
            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')


@app.route('/change_model', methods=['POST'])
def change_model():
    global model, model2, model2_name
    data = request.json
    model_type = data['model_type']
    model_size = data['model_size']

    # 根据模型类型和大小更新模型路径
    if model_type == "目标追踪":
        model_path = f"./weights/yolov8{model_size}.pt"
        if model_size == 's':
            model2_path = './weights/detection_module.pt'  # 烟雾明火
            model2_name = 'detection_module'
        else:
            model2_path = None
            model2_name = None
    elif model_type == "吸烟检测":
        model_path = "./weights/detection_module.pt"  # 吸烟检测
        model2_path = None
        model2_name = None
    elif model_type == "烟雾明火检测":
        model_path = './weights/best1.pt'  # 烟雾明火
        model2_path = None
        model2_name = None
    else:
        return jsonify({'error': 'Invalid model type'}), 400

    # 初始化模型
    model = init_model(model_path)
    if model2_path:
        model2 = init_model(model2_path)
    else:
        model2 = None

    model2_name = model2_name if model2_path else None

    logging.info(f"更改模型为 {model_path}")

    return jsonify({'status': 'Model changed'}), 200

@app.route('/pause_video', methods=['POST'])
def pause_video():
    global is_paused
    is_paused = True  # 暂停视频
    return jsonify({'status': 'Video paused'}), 200

@app.route('/play_video', methods=['POST'])
def play_video():
    global is_paused
    is_paused = False  # 恢复视频播放
    return jsonify({'status': 'Video playing'}), 200

@app.route('/replay_video', methods=['POST'])
def replay_video():
    global frame_position, cap
    frame_position = 0  # 重置到视频的起始位置
    cap.set(cv2.CAP_PROP_POS_FRAMES, frame_position)
    return jsonify({'status': 'Video replayed'}), 200

@app.route('/set_video_position', methods=['POST'])
def set_video_position():
    global cap, frame_position, is_paused
    data = request.json
    frame_position = data['frame_position']
    is_paused = True  # 设置视频位置时暂停视频
    if cap is not None:
        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_position)  # 设置视频位置
        is_paused = False  # 设置完成后恢复播放
        return jsonify({'status': 'Video position set'}), 200
    else:
        return jsonify({'error': 'Video capture is not initialized'}), 400

@app.route('/set_region', methods=['POST'])
def set_region():
    global region
    data = request.json
    try:
        region = [int(data['x1']), int(data['y1']), int(data['x2']), int(data['y2'])]
        return jsonify({'status': 'Region updated'}), 200
    except (KeyError, ValueError) as e:
        return jsonify({'error': str(e)}), 400

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5001)
